/*
 * dur_hmm_markov.cc
 *
 *  Created on: Sep 20, 2011
 *      Author: tqlong
 */

#include <boost/foreach.hpp>
#include "dur_hmm_markov.h"

void DurationHMMMarkov::load(const char* fileName)
{
  int n_state, n_observation;
  std::ifstream file(fileName);
  if (!file.good())
  {
    std::cerr << "Cannot open file " << fileName << std::endl;
    exit(1);
  }

  file >> n_state >> n_observation;
  start_ = arma::vec(n_state);
  trans_ = arma::mat(n_state, n_state);           // column wise p(i|j)
  emis_  = arma::mat(n_observation, n_state);     // column wise p(o=k|j)
  for (int i = 0; i < n_state; i++)
    file >> start_(i);
  for (int j = 0; j < n_state; j++)
    for (int i = 0; i < n_state; i++)
      file >> trans_(i, j);
  for (int j = 0; j < n_state; j++)
    for (int k = 0; k < n_observation; k++)
      file >> emis_(k, j);

  for (int i = 0; i < n_state; i++)
  {
    int n_dur;
    file >> n_dur;
    dprob_.push_back(duration_prob_t());
    for (int d = 0; d < n_dur; d++)
    {
      int d;
      double prob;
      file >> d >> prob;
      dprob_.back()[d] = prob;
    }
  }

  m_emis_.resize(n_state);
  for (int i = 0; i < n_state; i++)
  {
    m_emis_[i] = arma::mat(n_observation, n_observation);
    for (int kj = 0; kj < n_observation; kj++)
      for (int ki = 0; ki < n_observation; ki++)
        file >> m_emis_[i](ki, kj);
  }

  file.close();
}

void DurationHMMMarkov::save(const char* fileName)
{
  std::ofstream file(fileName);
  if (!file.good())
  {
    std::cerr << "Cannot open file " << fileName << std::endl;
    exit(1);
  }

  file << n_state() << " " << n_observation() << std::endl;
  for (int i = 0; i < n_state(); i++)
    file << start_(i) << " ";
  file << std::endl;
  for (int j = 0; j < n_state(); j++)
  {
    for (int i = 0; i < n_state(); i++)
      file << trans_(i, j) << " ";
    file << std::endl;
  }
  for (int j = 0; j < n_state(); j++)
  {
    for (int k = 0; k < n_observation(); k++)
      file << emis_(k, j) << " ";
    file << std::endl;
  }

  for (int i = 0; i < n_state(); i++)
  {
    int n_dur = dprob_[i].size();
    file << n_dur << " ";
    BOOST_FOREACH(const_pair_dprob_t p, dprob_[i])
    {
      file << p.first << " " << p.second << " ";
    }
    file << std::endl;
  }

  for (int i = 0; i < n_state(); i++)
  {
    for (int kj = 0; kj < n_observation(); kj++)
    {
      for (int ki = 0; ki < n_observation(); ki++)
        file << m_emis_[i](ki, kj) << " ";
      file << std::endl;
    }
  }

  file.close();
}

// randomization
void DurationHMMMarkov::rand()
{
  DurationHMM::rand();
  for (int i = 0; i < n_state(); i++)
  {
    m_emis_[i].randu(n_observation(), n_observation());
    for (int kj = 0; kj < n_observation(); kj++)
      m_emis_[i].col(kj) /= arma::sum(m_emis_[i].col(kj));
  }
}

void DurationHMMMarkov::generate(int T, joint_seq_t &seq)
{
  int state = discrete_sample(start_);
  sequence_t &o = seq.oseq, &q = seq.qseq;
  o.clear();
  q.clear();
  for (int t = 0; t < T; t++)
  {
    int d = duration_sample(dprob_[state]);
    int obs = discrete_sample(emis_.col(state));
    for (int dd = 0; dd < d; dd++)
    {
      q.push_back(state);
      o.push_back(obs);
      obs = discrete_sample(m_emis_[state].col(obs));
    }
    state = discrete_sample(trans_.col(state));
  }
}

// computation functions
// computation functions
double DurationHMMMarkov::seq_probability(const sequence_t& o)
{
  compute_alpha(o, var_);
  int T = o.size();
  double s = log(0);
  for (int i = 0; i < n_state(); i++)
    s = logadd(s, var_.alpha(i, T-1));

//  //tests for other functions
//  compute_beta(o, var_);
//  {
//    compute_alphaS(o, var_);
//    double ss = log(0);
//    for (int i = 0; i < n_state(); i++)
//      ss = logadd(ss, var_.alphaS(i, T));
//    std::cout << "prob(alphaS) = " << Log(ss) << std::endl;
//  }
//  {
//    compute_betaS(o, var_);
//    double ss = log(0);
//    for (int i = 0; i < n_state(); i++)
//      ss = logadd(ss, var_.betaS(i, 0)+log(start_(i)));
//    std::cout << "prob(betaS) = " << Log(ss) << std::endl;
//  }

  return s;
}

double DurationHMMMarkov::seq_probability_reverse(const sequence_t& o)
{
  compute_beta(o, var_);
  double s = log(0);
  for (int i = 0; i < n_state(); i++)
  {
    BOOST_FOREACH(const_pair_dprob_t p, dprob_[i])
    {
      int d = p.first;
      double prob = p.second;
      double tmp = log(start_(i)) + log(prob) + subseq_prob(o, 0, d, i) + var_.beta(i, d-1);
      s = logadd(s, tmp);
    }
  }
  return s;
}

void DurationHMMMarkov::map(joint_seq_t& seq)
{
  const sequence_t& o = seq.oseq;
  int T = o.size();
  var_.alpha = arma::mat(n_state(), T);

  // viterbi
  for (int t = 0; t < T; t++)
  {
    for (int i = 0; i < n_state(); i++)
    {
      double s = log(0);
      if (dprob_[i].find(t+1) != dprob_[i].end())                         // first possibility: state i starts at 0 and ends at t
        s = log(start_(i)) + log(dprob_[i][t+1]) + subseq_prob(o, 0, t+1, i);
      BOOST_FOREACH(const_pair_dprob_t p, dprob_[i])
      {
        int d = p.first;         // d
        double prob = p.second;  // p_i(d)
        if (t >= d)
        {
          double subseq_p = log(prob) + subseq_prob(o, t-d+1, d, i);
          for (int j = 0; j < n_state(); j++) if (i != j)
          {
            double tmp = var_.alpha(j, t-d) + log(trans_(i,j)) + subseq_p;
            if (s < tmp) s = tmp;
          }
        }
      }
      var_.alpha(i, t) = s;
    }
  }

  // trace back
  sequence_t& q = seq.qseq;
  q.resize(T);
  int max_i = 0, max_d;
  for (int i = 0; i < n_state(); i++)
    if (var_.alpha(i,T-1) > var_.alpha(max_i, T-1)) max_i = i;

  int t = T-1;
  for (; t >= 0; )
  {
    int next_i = max_i;
    double max = log(0);
    max_i = -1;
    max_d = -1;

    if (dprob_[next_i].find(t+1) != dprob_[next_i].end())
    {
      max_i = -1;
      max_d = t+1;
      max = log(start_(next_i)) + log(dprob_[next_i][t+1]) + subseq_prob(o, 0, t+1, next_i);
    }

    BOOST_FOREACH(const_pair_dprob_t p, dprob_[next_i])
    {
      int d = p.first;
      double prob = p.second;
      if (t >= d)
      {
        double subseq_p = log(prob) + subseq_prob(o, t-d+1, d, next_i);
        for (int i = 0; i < n_state(); i++)
        {
          double tmp = var_.alpha(i, t-d) + log(trans_(next_i,i)) + subseq_p;
          if (max < tmp)
          {
            max = tmp;
            max_i = i;
            max_d = d;
          }
        }
      }
    }

    for (int d = 0; d < max_d; d++)
    {
      q[t--] = next_i;
    }
  }
}

void DurationHMMMarkov::train(const vseq_t& data)
{
  var_.ACC_start = arma::vec(n_state());
  var_.ACC_trans = arma::mat(n_state(), n_state());
  var_.ACC_emis = arma::mat(n_observation(), n_state());
  var_.ACC_dprob = dprob_;
  var_.ACC_m_emis = m_emis_;

  int iter = 0;
  for (; iter < 50; iter++)
  {
    std::stringstream s;
    s << prefix_ << iter << ".txt";
    save(s.str().c_str());

    s.clear();

    var_.ACC_start.fill(log(0));
    var_.ACC_trans.fill(log(0));
    var_.ACC_emis.fill(log(0));
    for (int i = 0; i < n_state(); i++)
    {
      BOOST_FOREACH(pair_dprob_t p, var_.ACC_dprob[i])
      {
        p.second = log(0);
      }
      var_.ACC_m_emis[i].fill(log(0));
    }

    double likelihood = 0;
    for (unsigned int n = 0; n < data.size(); n++)
    {
      const sequence_t& o = data[n];
      double log_prob = seq_probability(o); // forward
//      std::cout << "-----alpha " << std::endl; print_params();
      compute_beta(o, var_);                // backward
//      std::cout << "-----beta " << std::endl; print_params();
      compute_alphaS(o, var_);
//      std::cout << "-----alphaS " << std::endl; print_params();
      compute_betaS(o, var_);
//      std::cout << "-----betaS " << std::endl; print_params();
      accumulate(o, -log_prob);
//      std::cout << "-----accu " << std::endl; print_params();
      likelihood += log_prob;
    }
    std::cout << "iter = " << iter << " likelihood = " << likelihood << std::endl;

    std::stringstream sf;
    sf << prefix_ << "likelihood.txt";
    std::ofstream llfile(sf.str().c_str(), std::ios::app);
    llfile << iter << " " << likelihood << std::endl;
    llfile.close();

//    std::cout << "-----before reestimate " << std::endl; print_params();
    reestimate();
  }

  std::stringstream s;
  s << prefix_ << iter << ".txt";
  save(s.str().c_str());
}

// print functions
void DurationHMMMarkov::print_params() const
{
  DurationHMM::print_params();
  for (int i = 0; i < n_state(); i++)
  {
    std::cout << "m_emis(" << i << ")=" << std::endl << m_emis_[i];
  }
}

// protected members
void DurationHMMMarkov::compute_alpha(const sequence_t& o, variable_t& var)
{
  int T = o.size();
  var.alpha = arma::mat(n_state(), T);
  for (int t = 0; t < T; t++)
  {
    for (int i = 0; i < n_state(); i++)
    {
      double s = log(0);
      if (dprob_[i].find(t+1) != dprob_[i].end())                         // first possibility: state i starts at 0 and ends at t
        s = log(start_(i)) + log(dprob_[i][t+1]) + subseq_prob(o, 0, t+1, i);
      BOOST_FOREACH(const_pair_dprob_t p, dprob_[i])
      {
        int d = p.first;         // d
        double prob = p.second;  // p_i(d)
        if (t >= d)
        {
          double subseq_p = log(prob) + subseq_prob(o, t-d+1, d, i);
          for (int j = 0; j < n_state(); j++) if (i != j)
          {
            double tmp = var.alpha(j, t-d) + log(trans_(i,j)) + subseq_p;
            s = logadd(s, tmp);
          }
        }
      }
      var.alpha(i, t) = s;
    }
  }
}

void DurationHMMMarkov::compute_alphaS(const sequence_t& o, variable_t& var)
{
  int T = o.size();
  var.alphaS = arma::mat(n_state(), T+1);
  for (int t = -1; t < T; t++)
  {
    for (int i = 0; i < n_state(); i++)
    {
      if (t == -1)
        var.alphaS(i,t+1) = log(start_(i));
      else
      {
        double s = log(0);
        for (int j = 0; j < n_state(); j++) if (i != j)
        {
          double tmp = var.alpha(j, t) + log(trans_(i,j));
          s = logadd(s, tmp);
        }
        var.alphaS(i,t+1) = s;
      }
    }
  }
}

void DurationHMMMarkov::compute_beta(const sequence_t& o, variable_t& var)
{
  int T = o.size();
  var.beta = arma::mat(n_state(), T);
  for (int i = 0; i < n_state(); i++)
    var.beta(i, T-1) = log(1);

  for (int t = T-2; t >= 0; t--)
  {
    for (int j = 0; j < n_state(); j++)
    {
      double s = log(0);
      for (int i = 0; i < n_state(); i++) if (i != j)
      {
        BOOST_FOREACH(const_pair_dprob_t p, dprob_[i])
        {
          int d = p.first;
          double prob = p.second;
          if (t + d < T)
          {
            double tmp = var.beta(i, t+d) + log(trans_(i,j)) + log(prob) + subseq_prob(o, t+1, d, i);
            s = logadd(s, tmp);
          }
        }
      }
      var.beta(j,t) = s;
    }
  }
}

void DurationHMMMarkov::compute_betaS(const sequence_t& o, variable_t& var)
{
  int T = o.size();
  var.betaS = arma::mat(n_state(), T+1);

  for (int t = -1; t < T; t++)
  {
    for (int j = 0; j < n_state(); j++)
    {
      if (t == T-1)
      {
        var.betaS(j,t+1) = log(1);
        continue;
      }
      double s = log(0);
      BOOST_FOREACH(const_pair_dprob_t p, dprob_[j])
      {
        int d = p.first;
        double prob = p.second;
        if (t + d < T)
        {
          double tmp = var.beta(j, t+d) + log(prob) + subseq_prob(o, t+1, d, j);
          s = logadd(s, tmp);
        }
      }
      var.betaS(j,t+1) = s;
    }
  }
}

double DurationHMMMarkov::diff(DurationHMMMarkov& hmm)
{
  double s = 0;
  s += arma::dot(start_-hmm.start_, start_-hmm.start_);
  s += arma::dot(trans_-hmm.trans_, trans_-hmm.trans_);
  s += arma::dot(emis_-hmm.emis_, emis_-hmm.emis_);
  for (int i = 0; i < n_state(); i++)
  {
    BOOST_FOREACH(const_pair_dprob_t p, dprob_[i])
    {
      s += (p.second - hmm.dprob_[i][p.first])*(p.second - hmm.dprob_[i][p.first]);
    }
  }
  for (int i = 0; i < n_state(); i++)
  {
    s += arma::dot(m_emis_[i]-hmm.m_emis_[i], m_emis_[i]-hmm.m_emis_[i]);
  }
  return ::sqrt(s);
}

double DurationHMMMarkov::subseq_prob(const sequence_t& o, int start, int len, int state)
{
  double s = log(emis_(o[start], state));
  int t = start+1;
  for (int d = 1; d < len; d++, t++)
  {
    s += log(m_emis_[state](o[t], o[t-1]));
  }
  return s;
}

void DurationHMMMarkov::accumulate(const sequence_t& o, double weight)
{
  int T = var_.alpha.n_cols;

  for (int j = 0; j < n_state(); j++)
    for (int i = 0; i < n_state(); i++) if (i != j)
    {
      double s = log(0);
      for (int t = 0; t < T-1; t++)
      {
        double tmp = var_.alpha(j,t) + log(trans_(i,j)) + var_.betaS(i,t+1);
        s = logadd(s, tmp);
      }
      s += weight;
      var_.ACC_trans(i,j) = logadd(var_.ACC_trans(i,j), s);
    }

  arma::vec s_begin(n_state()), s_end(n_state());
  s_begin.fill(log(0));
  s_end.fill(log(0));
  for (int t = 0; t < T-1; t++)
  {
    int kj = o[t], ki = o[t+1];
    for (int i = 0; i < n_state(); i++)
    {
      double s = var_.alphaS(i,t) + var_.betaS(i, t) + weight;
      s_begin(i) = logadd(s_begin(i), s);
      s = var_.alpha(i,t) + var_.beta(i, t) + weight;
      s_end(i) = logadd(s_end(i), s);

      double tmp = logsub(s_begin(i), s_end(i));
      var_.ACC_m_emis[i](ki, kj) = logadd(var_.ACC_m_emis[i](ki, kj), tmp);
    }
  }

  for (int t = 0; t < T; t++)
  {
    int k = o[t];
    for (int i = 0; i < n_state(); i++)
    {
      double s = var_.alphaS(i,t) + var_.betaS(i, t) + weight;
      var_.ACC_emis(k, i) = logadd(var_.ACC_emis(k, i), s);
    }
  }

//  std::cout << "alphaS=" << std::endl << var_.alphaS;
//  std::cout << "beta=" << std::endl << var_.beta;
  for (int i = 0; i < n_state(); i++)
  {
    BOOST_FOREACH(const_pair_dprob_t p, dprob_[i])
    {
      int d = p.first;
      double prob = p.second;
      double s = log(0);
      for (int t = -1; t < T; t++) if (t+1 >= 0 && t+d < T)
      {
        double subseq_p = log(prob) + subseq_prob(o, t+1, d, i);
        double tmp = var_.alphaS(i, t+1) + subseq_p + var_.beta(i, t+d);
//        std::cout << "     - " << i << " " << d << " "<< t << " = " << Log(subseq_p) << " " << Log(tmp) << std::endl;
        s = logadd(s, tmp);
      }
      s += weight;
//      std::cout << "   + " << i << " " << d << " = " << Log(s) << " " << prob << std::endl;
      var_.ACC_dprob[i][d] = logadd(var_.ACC_dprob[i][d], s);
    }
  }

  for (int i = 0; i < n_state(); i++)
  {
    double s = log(start_(i)) + var_.betaS(i,0) + weight;
    var_.ACC_start(i) = logadd(var_.ACC_start(i), s);
  }
}

void DurationHMMMarkov::reestimate()
{
  for (int j = 0; j < n_state(); j++)
  {
    double s = logsum_vec(var_.ACC_trans.col(j));
    for (int i = 0; i < n_state(); i++)
      trans_(i,j) = exp(var_.ACC_trans(i, j) - s);
  }

  for (int j = 0; j < n_state(); j++)
  {
    for (int kj = 0; kj < n_observation(); kj++)
    {
      double s = logsum_vec(var_.ACC_m_emis[j].col(kj));
      for (int ki = 0; ki < n_observation(); ki++)
        m_emis_[j](ki,kj) = exp(var_.ACC_m_emis[j](ki, kj) - s);
    }
  }

  for (int j = 0; j < n_state(); j++)
  {
    double s = logsum_vec(var_.ACC_emis.col(j));
    for (int k = 0; k < n_observation(); k++)
      emis_(k,j) = exp(var_.ACC_emis(k, j) - s);
  }

  for (int i = 0; i < n_state(); i++)
  {
    double s = log(0);
//    std::cout << "ACC_prob(" << i << "): ";
    BOOST_FOREACH(const_pair_dprob_t p, var_.ACC_dprob[i])
    {
      double prob = p.second;
      s = logadd(s, prob);
//      std::cout << Log(prob) << " ";
    }
//    std::cout << std::endl;

    BOOST_FOREACH(pair_dprob_t p, dprob_[i])
    {
      p.second = exp(var_.ACC_dprob[i][p.first] - s);
    }
  }

  double s = logsum_vec(var_.ACC_start);
  for (int i = 0; i < n_state(); i++)
    start_(i) = exp(var_.ACC_start(i) - s);
}
